from urllib.parse import quote_plus
from pymongo import MongoClient
import pandas as pd
def get_database_info(database_name,collection_name):
    # 数据表定义

    columns = ['companyID', 'year', 'materialCount', 'materialSize']
    df = pd.DataFrame(columns=columns)

    # MongoDB 连接信息
    mongo_host = 'xxxxx'
    mongo_port = 3717
    mongo_user = 'xxx'
    mongo_password = 'xxxx'
    escaped_password = quote_plus(mongo_password)
    escaped_host = quote_plus(mongo_host)

    # 连接 MongoDB
    connection_string = f'mongodb://{mongo_user}:{escaped_password}@{mongo_host}:{mongo_port}/admin'
    client = MongoClient(connection_string)

    # 指定要获取信息的数据库
    db = client[database_name]

    # 获取数据库中所有集合
    collection = db[collection_name]

    # 从文本文件中读取 companyId
    company_ids = []
    with open('companyId.txt', 'r') as file:
        company_ids = [line.strip() for line in file]
        # 聚合查询
    pipeline = [
        {
            "$addFields": {
                "year": {
                    "$year": {
                        "$dateFromString": {
                            "dateString": "$ctime",
                            "format": "%Y-%m-%d %H:%M:%S"
                        }
                    }
                }
            }
        },
        {
            "$match": {
                "companyId": {"$in": company_ids}
            }
        },
        {
            "$group": {
                "_id": {"year": "$year", "companyId": "$companyId"},
                "materialCount": {"$sum": 1},
                "materialSize": {"$sum": "$size"}
            }
        },
        {
            "$sort": {"_id.companyId": 1, "_id.year": 1}
        },
        {
            "$project": {
                "_id": 0,
                "companyID": "$_id.companyId",
                "year": "$_id.year",
                "materialCount": 1,
                "materialSize": 1
            }
        }
    ]

    results = collection.aggregate(pipeline)
    for result in results:
     try:
       df = df.append({'companyID': result['companyID'], 'year': result['year'], 'materialCount': result['materialCount'], 'materialSize': result['materialSize']}, ignore_index=True)
     except Exception as op_err:
       print(f"exception due to:  op_err")
    df.to_csv(f"material.csv", index=False)

if __name__ == "__main__":
    # 指定要获取信息的数据库名称
    database_name = "rms-resource"
    collection_name = "material20240605"
    get_database_info(database_name, collection_name)



